# -*- coding: utf-8 -*-
"""
文件名称: epoch_compare_plots.py

功能概述:
    基于各子实验的 step_metrics_summary.json, 构造 "按 epoch" 的对比:
        1. 柱状图:
           - 横轴: 不同子实验 (不同 "诊断 vs 其他题型" 比例)
           - 纵轴: 各指标在第 1 个 epoch 结束时的取值
           - 每个指标一张柱状图
        2. 折线图:
           - 横轴: epoch 序号 (1..max_epoch)
           - 纵轴: 对应指标在各 epoch 结束时的取值
           - 一张图中多条折线, 每条线对应一个子实验 (不同采样比例)
           - 若某子实验仅 1 个 epoch, 则该线只有一个点, 并在该 y 值画一条水平虚线作为参考线

    epoch 边界确定方法:
        - 对每个子实验:
          * 从 scored/step_metrics_summary.json 中获取所有 step (int)
          * 找到最大 step: max_step
          * 根据目录名解析该模型总 epoch 数 E:
              sub_exp_1_diag_only  -> 7
              sub_exp_1_p5         -> 6
              sub_exp_1_p20        -> 4
              sub_exp_1_p50        -> 2
              其他 (例如 sub_exp_1) -> 1
          * gap = max_step / E
          * 对于第 k 个 epoch (k=1..E):
              target_step = k * gap
              在所有 step 中寻找距离 target_step 最近的 step, 作为该 epoch 的代表 step

输出:
    - out_dir/epoch_metrics_summary.json  (各实验、各 epoch 的指标汇总)
    - out_dir/<metric>_epoch1_bar.png     (第 1 个 epoch 的柱状图)
    - out_dir/<metric>_epoch_line.png     (按 epoch 变化的折线图)

使用示例:
    python fundus_reasoner/eval/epoch_compare_plots.py \
        --eval-root experiments/eval_output \
        --out-dir experiments/eval_output/epoch_compare
"""

import argparse
import json
import logging
from pathlib import Path
from typing import Dict, Any, List, Tuple

import matplotlib
import matplotlib.pyplot as plt

# 全局字体路径: 默认使用当前工作目录下的本地字体文件
FONT_PATH = Path("./SarasaMonoCL-Regular.ttf")

# 折线/柱状图的样式参数
MARKER_SIZE = 8
LINE_WIDTH = 2
BAR_WIDTH = 0.6


logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 字体设置: 使用本地 ./SarasaMonoCL-Regular.ttf
# -----------------------------
def setup_font() -> None:
    """使用当前目录下的本地字体文件 ./SarasaMonoCL-Regular.ttf

    要求:
        - FONT_PATH 指向实际存在的 .ttf 文件
        - 若文件不存在或加载失败则抛异常, 防止悄悄退回系统默认字体导致中文渲染异常
    """
    font_file = FONT_PATH

    if not font_file.is_file():
        raise FileNotFoundError(
            f"字体文件不存在: {font_file} "
            f"（请确保 SarasaMonoCL-Regular.ttf 放在当前工作目录, "
            f"或修改 epoch_compare_plots.py 中的 FONT_PATH）"
        )

    try:
        import matplotlib.font_manager as fm

        # 显式注册该字体文件
        fm.fontManager.addfont(str(font_file))
        # 通过文件获取字体名称, 避免硬编码 family 名字
        font_name = fm.FontProperties(fname=str(font_file)).get_name()

        matplotlib.rcParams["font.family"] = font_name
        matplotlib.rcParams["axes.unicode_minus"] = False

        logger.info("成功加载本地字体: %s, 字体名: %s", font_file, font_name)
    except Exception as e:
        raise RuntimeError(
            f"加载字体文件失败: {font_file}, error={e}"
        ) from e


# -----------------------------
# 工具函数: 根据子目录名解析 epoch 数
# -----------------------------
def parse_epochs_from_name(exp_name: str) -> int:
    """根据子实验目录名解析该模型训练的 epoch 数

    约定:
        - 以 "_diag_only" 结尾 -> 7
        - 以 "_p5" 结尾       -> 6
        - 以 "_p20" 结尾      -> 4
        - 以 "_p50" 结尾      -> 2
        - 其他                -> 1  (例如 sub_exp_1)
    """
    if exp_name.endswith("_diag_only"):
        return 7
    if exp_name.endswith("_p5"):
        return 6
    if exp_name.endswith("_p20"):
        return 4
    if exp_name.endswith("_p50"):
        return 2
    # 默认: 全题型 (all)
    return 1


# -----------------------------
# 工具函数: 解析子目录名 -> ratio, display_name
# -----------------------------
def parse_ratio_display_name(exp_name: str) -> Tuple[int, str]:
    """根据子目录名解析 “其他题型保留比例” ratio 及中文显示名

    与前面脚本保持一致:
        - *_diag_only -> ratio=0, "仅诊断题 (其他题型 0%)"
        - *_pXX       -> ratio=XX, "诊断题 + 其他题型 XX%"
        - 其他        -> ratio=100, "诊断题 + 其他题型 100%"
    """
    ratio: int = 100

    if exp_name.endswith("_diag_only"):
        ratio = 0
    else:
        parts = exp_name.split("_")
        last = parts[-1]
        if last.startswith("p") and last[1:].isdigit():
            ratio = int(last[1:])
        else:
            ratio = 100

    if ratio == 0:
        display = "仅诊断题 (其他题型 0%)"
    else:
        display = f"诊断题 + 其他题型 {ratio}%"

    return ratio, display


# -----------------------------
# 工具函数: 读取 summary JSON
# -----------------------------
def load_experiment_summary(exp_dir: Path) -> Dict[int, Dict[str, Any]]:
    """从子实验目录中读取 scored/step_metrics_summary.json

    返回:
        { step_int: metrics(dict), ... }
    """
    summary_path = exp_dir / "scored" / "step_metrics_summary.json"
    if not summary_path.is_file():
        raise RuntimeError(
            f"子实验目录缺少 scored/step_metrics_summary.json: {summary_path}"
        )

    try:
        raw = json.loads(summary_path.read_text(encoding="utf-8"))
    except Exception as e:
        raise RuntimeError(
            f"读取或解析子实验 summary JSON 失败: path={summary_path}, error={e}"
        ) from e

    if not isinstance(raw, dict):
        raise RuntimeError(
            f"summary JSON 顶层结构必须为 dict (step->metrics), 实际为: {type(raw)}"
        )

    step_map: Dict[int, Dict[str, Any]] = {}
    for k, v in raw.items():
        try:
            step_int = int(k)
        except Exception:
            raise RuntimeError(
                f"summary 中存在无法转换为 int 的 step 字符串: {k!r}"
            )
        if not isinstance(v, dict):
            raise RuntimeError(
                f"step={k} 对应条目不是 dict 类型, 实际为: {type(v)}"
            )
        step_map[step_int] = v

    if not step_map:
        raise RuntimeError(f"summary 中未发现任何 step 记录: {summary_path}")

    return step_map


# -----------------------------
# 工具函数: 给定 step->metrics, 根据 epoch 数切分 epoch 边界
# -----------------------------
def compute_epoch_metrics_for_experiment(
    steps_metrics: Dict[int, Dict[str, Any]],
    epochs: int,
) -> Tuple[Dict[int, int], Dict[int, Dict[str, float]]]:
    """针对单个实验, 根据最大 step 和 epochs 计算每个 epoch 对应的 step 与指标

    参数:
        steps_metrics: { step_int: metrics_dict } (来自 summary)
        epochs:        该实验训练的总 epoch 数

    返回:
        epoch_to_step:  { epoch_idx: chosen_step_int }
        epoch_metrics:  { epoch_idx: {metric_name: float, ...} }

    策略:
        - max_step = max(steps)
        - gap = max_step / epochs
        - 对于 epoch=k (1..epochs):
            target = k * gap
            在所有 steps 中寻找 abs(step - target) 最小的 step 作为 chosen_step
    """
    steps = sorted(steps_metrics.keys())
    max_step = max(steps)
    gap = max_step / float(epochs)

    epoch_to_step: Dict[int, int] = {}
    epoch_metrics: Dict[int, Dict[str, float]] = {}

    metric_keys = [
        "overall_strict_accuracy",
        "judge_strict_accuracy",
        "compare_strict_accuracy",
        "diagnosis_strict_accuracy",
        "choice_strict_accuracy",
    ]

    for epoch_idx in range(1, epochs + 1):
        target = epoch_idx * gap

        # 找最接近 target 的 step
        best_step = None
        best_dist = None
        for s in steps:
            dist = abs(s - target)
            if (best_dist is None) or (dist < best_dist):
                best_dist = dist
                best_step = s

        if best_step is None:
            raise RuntimeError(
                f"无法为 epoch={epoch_idx} 找到对应 step (steps 为空?)"
            )

        epoch_to_step[epoch_idx] = best_step

        metrics_raw = steps_metrics[best_step]
        metrics = {
            mk: float(metrics_raw.get(mk, 0.0)) for mk in metric_keys
        }
        epoch_metrics[epoch_idx] = metrics

    return epoch_to_step, epoch_metrics


# -----------------------------
# 总汇总: 读取所有子实验, 计算 epoch metrics
# -----------------------------
def gather_epoch_metrics(
    eval_root: Path,
) -> Dict[str, Dict[str, Any]]:
    """在 eval_root 下汇总所有子实验的 epoch 级指标

    返回结构:
        {
          exp_name: {
            "epochs": int,
            "ratio": int,
            "display_name": str,
            "epoch_to_step": {epoch_idx: step_int},
            "epoch_metrics": {epoch_idx: {metric_name: float, ...}},
          },
          ...
        }
    """
    if not eval_root.is_dir():
        raise RuntimeError(f"评估根目录不存在或不是目录: {eval_root}")

    subdirs = [d for d in eval_root.iterdir() if d.is_dir()]
    if not subdirs:
        raise RuntimeError(f"评估根目录下未发现任何子实验子目录: {eval_root}")

    logger.info("在评估根目录 %s 下发现 %d 个子实验目录", eval_root, len(subdirs))

    all_result: Dict[str, Dict[str, Any]] = {}

    for d in subdirs:
        exp_name = d.name
        epochs = parse_epochs_from_name(exp_name)
        ratio, display_name = parse_ratio_display_name(exp_name)

        steps_metrics = load_experiment_summary(d)
        epoch_to_step, epoch_metrics = compute_epoch_metrics_for_experiment(
            steps_metrics=steps_metrics,
            epochs=epochs,
        )

        all_result[exp_name] = {
            "epochs": epochs,
            "ratio": ratio,
            "display_name": display_name,
            "epoch_to_step": epoch_to_step,
            "epoch_metrics": epoch_metrics,
        }

        logger.info(
            "子实验 %s: epochs=%d, ratio=%d, max_step=%d",
            exp_name,
            epochs,
            ratio,
            max(steps_metrics.keys()),
        )

    return all_result


# -----------------------------
# 画柱状图: 第 1 个 epoch 对比
# -----------------------------
def plot_epoch1_bar_charts(
    epoch_result: Dict[str, Dict[str, Any]],
    out_dir: Path,
) -> None:
    """绘制第 1 个 epoch 结束后, 不同子实验的指标柱状图

    参数:
        epoch_result: gather_epoch_metrics 的返回值
        out_dir:      输出目录
    """
    out_dir.mkdir(parents=True, exist_ok=True)

    metric_keys = [
        "overall_strict_accuracy",
        "judge_strict_accuracy",
        "compare_strict_accuracy",
        "diagnosis_strict_accuracy",
        "choice_strict_accuracy",
    ]

    # 为了顺序美观, 按 ratio 排序
    exp_names_sorted = sorted(
        epoch_result.keys(),
        key=lambda name: epoch_result[name]["ratio"],
    )

    labels = [epoch_result[name]["display_name"] for name in exp_names_sorted]

    # 各指标一张图
    for mk in metric_keys:
        values: List[float] = []
        for name in exp_names_sorted:
            info = epoch_result[name]
            epoch_metrics = info["epoch_metrics"]
            # 第一个 epoch 的指标
            m = epoch_metrics[1][mk]
            values.append(m)

        plt.figure(figsize=(10, 6))
        ax = plt.gca()
        ax.set_facecolor("#f0f0f0")
        ax.grid(axis="y", linestyle="--", alpha=0.5)

        x_pos = list(range(len(exp_names_sorted)))
        ax.bar(
            x_pos,
            values,
            width=BAR_WIDTH,
        )

        ax.set_xticks(x_pos)
        ax.set_xticklabels(labels, rotation=20, ha="right")

        ax.set_ylabel("strict_accuracy", fontsize=12)
        ax.set_title(f"{mk} - 第 1 个 epoch 的模型对比", fontsize=14)

        plt.tight_layout()

        out_png = out_dir / f"{mk}_epoch1_bar.png"
        try:
            plt.savefig(out_png, dpi=600)
            logger.info("柱状图已保存: %s", out_png)
        except Exception as e:
            raise RuntimeError(f"保存柱状图失败: path={out_png}, error={e}") from e
        finally:
            plt.close()


# -----------------------------
# 画折线图: 按 epoch 的变化 (含单点虚线)
# -----------------------------
def plot_epoch_line_charts(
    epoch_result: Dict[str, Dict[str, Any]],
    out_dir: Path,
) -> None:
    """为每个指标画 epoch 维度的折线图

    参数:
        epoch_result: gather_epoch_metrics 的返回值
        out_dir:      输出目录

    说明:
        - X 轴: epoch 序号 (1..max_epochs)
        - Y 轴: 对应指标在各 epoch 结束时的值
        - 一张图中多条折线, 每条线对应一个子实验
        - 若某实验只有 1 个 epoch:
            * 在该 epoch 处画一个方形点
            * 并在该 y 值画一条横向虚线作为参考线
    """
    out_dir.mkdir(parents=True, exist_ok=True)

    metric_keys = [
        "overall_strict_accuracy",
        "judge_strict_accuracy",
        "compare_strict_accuracy",
        "diagnosis_strict_accuracy",
        "choice_strict_accuracy",
    ]

    # 按 ratio 排序决定线的颜色与图例顺序
    exp_names_sorted = sorted(
        epoch_result.keys(),
        key=lambda name: epoch_result[name]["ratio"],
    )

    # 使用 matplotlib 默认颜色循环
    color_cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", [])
    if not color_cycle:
        color_cycle = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]

    exp_colors: Dict[str, str] = {}
    for idx, name in enumerate(exp_names_sorted):
        exp_colors[name] = color_cycle[idx % len(color_cycle)]

    # 计算全局最大 epoch 以便设定 x 轴范围
    max_epoch_global = max(info["epochs"] for info in epoch_result.values())

    for mk in metric_keys:
        plt.figure(figsize=(10, 6))
        ax = plt.gca()
        ax.set_facecolor("#f0f0f0")
        ax.grid(True, linestyle="--", alpha=0.5)

        for name in exp_names_sorted:
            info = epoch_result[name]
            epochs = info["epochs"]
            metrics_map = info["epoch_metrics"]
            display_name = info["display_name"]
            color = exp_colors[name]

            x_vals: List[int] = []
            y_vals: List[float] = []

            for e in range(1, epochs + 1):
                m = metrics_map[e][mk]
                x_vals.append(e)
                y_vals.append(m)

            if epochs == 1:
                # 只有 1 个 epoch:
                # 1) 在该点画一个方形点(不进 legend)
                ax.plot(
                    x_vals,
                    y_vals,
                    marker="s",
                    linestyle="None",
                    markersize=MARKER_SIZE,
                    color=color,
                    label="_nolegend_",
                )
                # 2) 在该 y 值画一条横向虚线作为参考线 (带 legend)
                y0 = y_vals[0]
                ax.hlines(
                    y0,
                    xmin=1,
                    xmax=max_epoch_global,
                    colors=color,
                    linestyles="--",
                    linewidth=LINE_WIDTH,
                    label=f"{display_name} (epoch=1, 参考线)",
                )
            else:
                # 正常多 epoch 情况: 折线 + 坐标点
                ax.plot(
                    x_vals,
                    y_vals,
                    marker="s",
                    linestyle="-",
                    linewidth=LINE_WIDTH,
                    markersize=MARKER_SIZE,
                    color=color,
                    label=f"{display_name} (epoch={epochs})",
                )

        ax.set_xlabel("epoch", fontsize=12)
        ax.set_ylabel("strict_accuracy", fontsize=12)
        ax.set_title(f"{mk} 随 epoch 变化曲线", fontsize=14)

        ax.set_xlim(1, max_epoch_global)
        ax.set_xticks(range(1, max_epoch_global + 1))

        ax.legend(loc="best", fontsize=9)

        plt.tight_layout()

        out_png = out_dir / f"{mk}_epoch_line.png"
        try:
            plt.savefig(out_png, dpi=600)
            logger.info("折线图已保存: %s", out_png)
        except Exception as e:
            raise RuntimeError(f"保存折线图失败: path={out_png}, error={e}") from e
        finally:
            plt.close()


# -----------------------------
# 主逻辑
# -----------------------------
def run_epoch_compare(
    eval_root: Path,
    out_dir: Path,
) -> None:
    """主流程: 汇总 epoch 指标, 生成柱状图 + 折线图 + 汇总 JSON"""
    setup_font()

    epoch_result = gather_epoch_metrics(eval_root)

    out_dir.mkdir(parents=True, exist_ok=True)

    # 写一个汇总 JSON
    summary_obj = {}
    for name, info in epoch_result.items():
        summary_obj[name] = {
            "epochs": info["epochs"],
            "ratio": info["ratio"],
            "display_name": info["display_name"],
            "epoch_to_step": {
                str(e): int(s) for e, s in sorted(info["epoch_to_step"].items())
            },
            "epoch_metrics": {
                str(e): info["epoch_metrics"][e]
                for e in sorted(info["epoch_metrics"].keys())
            },
        }

    summary_path = out_dir / "epoch_metrics_summary.json"
    try:
        summary_path.write_text(
            json.dumps(summary_obj, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("epoch 级指标汇总 JSON 已写入: %s", summary_path)
    except Exception as e:
        raise RuntimeError(
            f"写入 epoch_metrics_summary.json 失败: path={summary_path}, error={e}"
        ) from e

    # 画柱状图 (第 1 个 epoch 对比)
    plot_epoch1_bar_charts(epoch_result, out_dir=out_dir)

    # 画折线图 (按 epoch 的变化)
    plot_epoch_line_charts(epoch_result, out_dir=out_dir)


# -----------------------------
# CLI
# -----------------------------
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="按 epoch 对比各子实验模型效果 (柱状图 + 折线图)",
    )
    parser.add_argument(
        "--eval-root",
        "-i",
        type=str,
        required=True,
        help="评估根目录, 其下每个子目录为一个子实验, 且含 scored/step_metrics_summary.json",
    )
    parser.add_argument(
        "--out-dir",
        "-o",
        type=str,
        default=None,
        help="输出目录 (默认: <eval-root>/epoch_compare)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    eval_root = Path(args.eval_root)
    if args.out_dir is not None:
        out_dir = Path(args.out_dir)
    else:
        out_dir = eval_root / "epoch_compare"

    run_epoch_compare(eval_root=eval_root, out_dir=out_dir)


if __name__ == "__main__":
    main()
